import os
# 完全禁用GPU以避免CUDA相关错误
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

import tkinter as tk
from tkinter import filedialog, ttk
from PIL import Image, ImageTk, ImageDraw, ImageFont
import numpy as np
import tensorflow as tf
import sys
import threading

# 添加项目根目录到Python路径
sys.path.append(os.path.abspath("."))

# 导入Faster R-CNN模型
from detection.models.detectors import faster_rcnn

classes = ['bg', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
           'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
           'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
           'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
           'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
           'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
           'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
           'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
           'hair drier', 'toothbrush']

# 推荐的模型加载方式
def load_faster_rcnn_model(num_classes, weights_path=None):
    """
    加载Faster R-CNN模型并处理权重
    
    Args:
        num_classes: 类别数量
        weights_path: 权重文件路径（可选）
    
    Returns:
        model: 加载好的模型实例
    """
    # 1. 创建模型实例
    model = faster_rcnn.FasterRCNN(num_classes=num_classes)

    # 2. 构建模型结构（通过一次前向传播）
    dummy_image = tf.zeros((1, 800, 800, 3))
    dummy_meta = [{
        'original_size': (tf.constant(800, dtype=tf.int32), tf.constant(800, dtype=tf.int32)),
        'target_size': (tf.constant(800, dtype=tf.int32), tf.constant(800, dtype=tf.int32)),
        'scale_x': tf.constant(1.0, dtype=tf.float32),
        'scale_y': tf.constant(1.0, dtype=tf.float32)
    }]

    # 调用模型以初始化所有层
    _ = model([dummy_image, dummy_meta], training=False)

    # 3. 如果有权重路径，则加载权重
    if weights_path and os.path.exists(weights_path):
        try:
            # 使用by_name=True和skip_mismatch=True来处理可能的不匹配
            model.load_weights(weights_path, by_name=True, skip_mismatch=True)
            print("权重加载成功")
        except Exception as e:
            print(f"权重加载失败: {e}")
            print("继续使用随机初始化的权重")
    else:
        if weights_path:
            print(f"权重文件不存在: {weights_path}")
        print("使用随机初始化的权重")

    return model

# 创建全局模型实例
model = load_faster_rcnn_model(num_classes=len(classes), weights_path='weights/faster_rcnn.h5')

def preprocess_image(image):
    """
    对图像进行预处理，转换为 Fast R-CNN 需要的格式。
    :param image: PIL Image 对象
    :return: 预处理后的图像数组和图像元信息
    """
    # 保存原始图像尺寸
    original_img = image.copy()
    original_size = image.size  # (width, height)
    
    # 调整图像尺寸，这里假设模型需要固定大小的输入，例如 800x800
    target_size = (800, 800)
    image = image.resize(target_size)
    
    # 转换为numpy数组并归一化处理
    image_array = np.array(image)
    image_array = image_array / 255.0
    
    # 计算缩放比例
    scale_x = target_size[0] / original_size[0]
    scale_y = target_size[1] / original_size[1]
    
    # 图像元信息，包含原始尺寸和缩放比例
    imagemeta = {
        'original_size': (tf.constant(original_size[1], dtype=tf.int32), tf.constant(original_size[0], dtype=tf.int32)),  # (height, width)
        'target_size': (tf.constant(target_size[1], dtype=tf.int32), tf.constant(target_size[0], dtype=tf.int32)),      # (height, width)
        'scale_x': tf.constant(scale_x, dtype=tf.float32),
        'scale_y': tf.constant(scale_y, dtype=tf.float32)
    }
    
    return image_array, imagemeta, original_img

def visualize_detections(image, boxes, scores, class_ids, scale_x, scale_y):
    """
    将检测结果可视化到图像上。
    :param image: PIL Image 对象
    :param boxes: 检测框坐标数组
    :param scores: 检测得分数组
    :param class_ids: 类别 ID 数组
    :param scale_x: x轴缩放比例
    :param scale_y: y轴缩放比例
    :return: 可视化后的 PIL Image 对象
    """
    draw = ImageDraw.Draw(image)
    try:
        # 尝试加载更好的字体
        font = ImageFont.truetype("arial.ttf", 15)
    except:
        # 如果加载失败，使用默认字体
        font = ImageFont.load_default()

    for i, (box, score, class_id) in enumerate(zip(boxes, scores, class_ids)):
        if score > 0.5:  # 过滤低置信度检测结果
            # box格式: [y1, x1, y2, x2]
            y1, x1, y2, x2 = box
            class_name = classes[class_id]
            
            # 将检测框坐标还原到原始图像尺寸
            x1_original = int(x1 / scale_x)
            y1_original = int(y1 / scale_y)
            x2_original = int(x2 / scale_x)
            y2_original = int(y2 / scale_y)

            # 绘制检测框
            draw.rectangle([x1_original, y1_original, x2_original, y2_original], outline="red", width=2)
            
            # 绘制标签文本
            text = f"{class_name}: {score:.2f}"
            # 获取文本尺寸
            try:
                text_width, text_height = font.getsize(text)
            except:
                text_width, text_height = (len(text) * 6, 12)  # 估算尺寸
            
            # 绘制文本背景
            draw.rectangle([x1_original, y1_original - text_height, x1_original + text_width, y1_original], 
                          fill="red")
            # 绘制文本
            draw.text((x1_original, y1_original - text_height), text, fill="white", font=font)
            
    return image

class ImageViewer:
    def __init__(self, root):
        self.root = root
        self.root.title("Faster R-CNN 目标检测")
        self.root.geometry("1200x700")  # 增加窗口尺寸以容纳更多组件

        # 创建主框架
        main_frame = tk.Frame(root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 创建按钮框架
        button_frame = tk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=(0, 10))
        
        # 创建按钮
        self.btn_select = tk.Button(button_frame, text="选择图像", command=self.select_image)
        self.btn_select.pack(side=tk.LEFT, padx=(0, 10))
        
        self.btn_detect = tk.Button(button_frame, text="开始检测", command=self.start_detection, state=tk.DISABLED)
        self.btn_detect.pack(side=tk.LEFT)
        
        # 创建进度条
        self.progress = ttk.Progressbar(button_frame, length=200, mode='indeterminate')
        self.progress.pack(side=tk.RIGHT, padx=(10, 0))
        tk.Label(button_frame, text="处理进度:").pack(side=tk.RIGHT)

        # 创建图像显示框架
        image_frame = tk.Frame(main_frame)
        image_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建原始图像框架
        original_frame = tk.LabelFrame(image_frame, text="原始图像", padx=5, pady=5)
        original_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=(0, 5))
        
        # 创建检测结果框架
        result_frame = tk.LabelFrame(image_frame, text="检测结果", padx=5, pady=5)
        result_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=(5, 0))
        
        # 创建图像显示标签
        self.original_label = tk.Label(original_frame)
        self.original_label.pack(expand=True)
        
        self.result_label = tk.Label(result_frame)
        self.result_label.pack(expand=True)
        
        # 存储图像数据
        self.original_img = None
        self.processed_img = None
        self.imagemeta = None
        self.original_img_copy = None
        self.scale_x = 1.0
        self.scale_y = 1.0

    def select_image(self):
        # 打开文件选择对话框
        file_path = filedialog.askopenfilename(
            title="选择图像文件",
            filetypes=[("图像文件", "*.png;*.jpg;*.jpeg;*.bmp")]
        )

        if not file_path:  # 用户取消选择
            return

        try:
            # 加载原始图像
            self.original_img = Image.open(file_path)
            
            # 保存一份副本用于显示
            self.original_img_copy = self.original_img.copy()
            
            # 预处理图像
            self.processed_img, self.imagemeta, _ = preprocess_image(self.original_img)
            self.scale_x = float(self.imagemeta['scale_x'])
            self.scale_y = float(self.imagemeta['scale_y'])
            
            # 显示原始图像
            self.display_original_image()
            
            # 启用检测按钮
            self.btn_detect.config(state=tk.NORMAL)
            
        except Exception as e:
            print(f"图像加载出错: {e}")
            import traceback
            traceback.print_exc()

    def display_original_image(self):
        """显示原始图像"""
        if self.original_img_copy:
            # 调整图像尺寸以适应显示
            display_img = self.original_img_copy.copy()
            display_img.thumbnail((500, 400))  # 限制最大显示尺寸
            self.tk_original_img = ImageTk.PhotoImage(display_img)
            self.original_label.config(image=self.tk_original_img)

    def start_detection(self):
        """在新线程中启动检测过程"""
        # 禁用按钮防止重复点击
        self.btn_select.config(state=tk.DISABLED)
        self.btn_detect.config(state=tk.DISABLED)
        
        # 启动进度条
        self.progress.start()
        
        # 在新线程中运行检测
        detection_thread = threading.Thread(target=self.run_detection)
        detection_thread.daemon = True
        detection_thread.start()

    def run_detection(self):
        """执行目标检测"""
        try:
            # 构造模型输入
            model_input = [
                tf.expand_dims(self.processed_img, 0),  # 图像数据 (1, 800, 800, 3)
                [self.imagemeta]  # 图像元信息
            ]
            
            # 调用模型进行推理
            results = model(model_input, training=False)
            
            # 处理模型输出
            if isinstance(results, list):
                # 如果返回的是列表，提取检测结果
                detections = results[0] if len(results) > 0 else None
            else:
                # 如果直接返回张量
                detections = results
            
            if detections is not None:
                # 转换为numpy数组
                detections_np = detections.numpy()
                
                # 提取检测框、得分和类别ID
                # detections格式: [batch, num_detections, (y1, x1, y2, x2, class_id, score)]
                boxes = detections_np[0, :, :4]      # (y1, x1, y2, x2)
                scores = detections_np[0, :, 5]      # score
                class_ids = detections_np[0, :, 4].astype(int)  # class_id
            else:
                print("模型未返回有效检测结果")
                boxes = np.array([])
                scores = np.array([])
                class_ids = np.array([])
                
            # 在主线程中更新UI
            self.root.after(0, self.update_detection_results, boxes, scores, class_ids)
                
        except Exception as e:
            print(f"目标检测出错: {e}")
            import traceback
            traceback.print_exc()
            # 在主线程中更新UI
            self.root.after(0, self.detection_finished)

    def update_detection_results(self, boxes, scores, class_ids):
        """更新检测结果"""
        try:
            # 可视化检测结果
            if len(boxes) > 0:
                result_img = visualize_detections(
                    self.original_img_copy.copy(), 
                    boxes, 
                    scores, 
                    class_ids, 
                    self.scale_x, 
                    self.scale_y
                )
            else:
                result_img = self.original_img_copy.copy()
                # 显示无检测结果的提示
                draw = ImageDraw.Draw(result_img)
                try:
                    font = ImageFont.truetype("arial.ttf", 30)
                except:
                    font = ImageFont.load_default()
                draw.text((10, 10), "未检测到目标", fill="red", font=font)
            
            # 调整图像尺寸以适应显示
            result_img.thumbnail((500, 400))  # 限制最大显示尺寸
            self.tk_result_img = ImageTk.PhotoImage(result_img)
            self.result_label.config(image=self.tk_result_img)
            
        except Exception as e:
            print(f"更新检测结果出错: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.detection_finished()

    def detection_finished(self):
        """检测完成后的清理工作"""
        # 停止进度条
        self.progress.stop()
        
        # 重新启用按钮
        self.btn_select.config(state=tk.NORMAL)
        if self.original_img is not None:
            self.btn_detect.config(state=tk.NORMAL)

if __name__ == "__main__":
    root = tk.Tk()
    app = ImageViewer(root)
    root.mainloop()